package nl.cypherpunk.statefuzzer;

import net.automatalib.words.Alphabet;
import nl.cypherpunk.statefuzzer.openvpn.*;

import java.io.*;
import java.net.SocketException;
import java.nio.charset.Charset;
import java.nio.file.Files;
import java.nio.file.LinkOption;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.nio.file.attribute.BasicFileAttributeView;
import java.nio.file.attribute.BasicFileAttributes;
import java.util.*;

public class VPNReplay {
    //paragrams parse
    //args[0] : tcp/udp
    //args[1] : auth
    //args[2] : cipher
    //args[3] : target command
    //args[4] : /home/pany/test/libgcov_preload.so
    //args[5] : gcovr command
    //args[6] : step to skip running gcovr and outputting data to covfile
    //          e.g., step=5 means we run gcovr after every 5 test cases
    //args[7] : starttime
    //args[8] : endtime
    //args[9] : seed directory
    public static void main(String[] args) throws Exception {
        Alphabet<String> alphabet;
        VPNTestService vpn;
        if(args[0].equals("tcp")) {
            vpn = new VPNTestServiceTCP();
        } else {
            vpn = new VPNTestServiceUDP();
        }
        vpn.setTarget("server");
        vpn.setLocal("127.0.0.1");
        vpn.setRemote("127.0.0.1");
        vpn.setLocalPort(33867);
        vpn.setRemotePort(1194);
        vpn.setAuth(args[1]);
        vpn.setCipher(args[2]);
        vpn.setCommand(args[3]);
        vpn.setEnvCmd(args[4]);
        try {
            vpn.start();
        } catch(SocketException e) {
            e.printStackTrace();
            vpn.closeSocket();
            vpn.close();
        }
        int step =  Integer.parseInt(args[6]);
        File file = new File(args[9]);
        File[] tempList = file.listFiles();
        List fileList = Arrays.asList(tempList);
        Collections.sort(fileList, new Comparator<File>() {
            @Override
            public int compare(File o1, File o2) {
                if (o1.isDirectory() && o2.isFile())
                    return -1;
                if (o1.isFile() && o2.isDirectory())
                    return 1;
                return o1.getName().compareTo(o2.getName());
            }
        });
        Long startTime = Long.parseLong(args[7]);
        Long endTime = Long.parseLong(args[8]);
        Long timestamp = 0L;
        Long delta = 0L;
        int totalNum = 0;
        for (File file1 : tempList) {
            File temp = new File(file1.getPath());
            System.out.println(file1.getName());
            totalNum++;
            try {
                try {
                    Path path = Paths.get(file1.getPath());
                    BasicFileAttributeView basicview = Files.getFileAttributeView(path, BasicFileAttributeView.class, LinkOption.NOFOLLOW_LINKS);
                    BasicFileAttributes attr = basicview.readAttributes();
                    timestamp = (Long)attr.creationTime().toMillis()/1000;

                } catch (Exception e) {
                    e.printStackTrace();
                    timestamp = (Long)file.lastModified()/1000;
                }
                if (totalNum == 1)
                {
                    delta = timestamp - startTime;
                }
                timestamp = timestamp - delta;
                InputStream fileIn = new FileInputStream(temp);
                //DataInputStream in = new DataInputStream(fileIn);

                // 使用缓存区读入对象效率更快
                BufferedInputStream in = new BufferedInputStream(fileIn);
                byte[] len = new byte[1024];
                int length;
                List<byte[]> seed = new ArrayList<>();
                vpn.reset();
                while (in.read(len, 0, 4) != -1) {
                    length = byteArrayToInt(len);
                    byte[] bys = new byte[length];
                    in.read(bys, 0, length);
                    try {
                        String output = vpn.sendFuzzingMessage(bys);
                        if (output.contains("Closed"))
                            break;
                    } catch (Exception e) {
                        System.out.println(e);
                    }
                    seed.add(bys);
                }
                in.close();
                fileIn.close();
            } catch (Exception e) {
                e.printStackTrace();
            }
            if(totalNum % step == 0 || totalNum == tempList.length || timestamp > endTime){
                List<String> cmdList = new ArrayList<>();
                cmdList.add("/bin/sh");
                cmdList.add("-cf");
                cmdList.add(args[5]);
                ProcessBuilder pb = new ProcessBuilder(cmdList);
                pb.redirectErrorStream(true);
                Process process = pb.start();
                int coverage = readProcessOutput(process.getInputStream(), "coverage.txt", timestamp);
                process.waitFor();
                process.exitValue();
            }
        }
        vpn.closeSocket();
        vpn.close();
        return;
    }

    private static Integer readProcessOutput(InputStream inputStream, String filename, Long seedTime) throws
            FileNotFoundException {
        PrintStream log = new PrintStream(new BufferedOutputStream(new FileOutputStream(filename, true)));

        Double[] coverage = new Double[4];
        try {
            BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream, Charset.forName("GBK")));
            String line;

            if ((line = reader.readLine()) != null) {
                //System.out.println(line);
                String[] tmp = line.split(" ");
                coverage[0] = Double.parseDouble(tmp[1].substring(0, tmp[1].length() - 1));
                coverage[1] = Double.parseDouble(tmp[2].substring(1));
                coverage[2] = Double.parseDouble(tmp[7].substring(0, tmp[7].length() - 1));
                coverage[3] = Double.parseDouble(tmp[8].substring(1));
            }
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                inputStream.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        log.print(seedTime);
        log.print(" ");
        log.print(coverage[0]);
        log.print(" ");
        log.print(coverage[1]);
        log.print(" ");
        log.print(coverage[2]);
        log.print(" ");
        log.println(coverage[3]);
        log.close();
        return coverage[3].intValue();
    }

    public static int byteArrayToInt(byte[] bytes) {
        int value=0;
        for(int i = 0; i < 4; i++) {
            int shift= i * 8;
            value +=(bytes[i] & 0xFF) << shift;
        }
        return value;
    }
}
